import glob
import os.path

from PIL import Image

import cv2
import numpy as np

save_path = "/home/rhq/xiangmu/yolov7-yy/2023-02-27-17-39-45_c25/training/mask"
img = cv2.imread("/home/rhq/xiangmu/mask.png")
mask = np.zeros((720, 1280, 3))
# black = np.array([0, 0, 0])
for i in range(0, 720):
    for j in range(0, 1280):
        if sum(img[i, j, 0:3]) != 0:
            mask[i, j, 0:3] = 1

img_paths = glob.glob("/home/rhq/xiangmu/yolov7-yy/8014/2023-02-27-17-39-45/camera/front/*.png")
for image_path in img_paths:
    img_name = image_path.split('.')[0].split('/')[-1]
    image = cv2.imread(image_path)
    new_img = image * mask
    save = os.path.join(save_path, img_name + '.png')
    cv2.imwrite(save, new_img)
